package leetCodeEveryday.Test2488统计中位数为K的子数组;

import java.util.HashMap;

public class Solution {
    //超出内存限制 气晕
    public int countSubarrays(int[] nums, int k) {
        //dp[i][j]: nums[i~j]中大于k的数字的数量减去小于k的数字的数量
        int[][] dp = new int[nums.length][nums.length];
        for (int i = 0; i < nums.length; i++) {
            for (int j = 0; j < nums.length; j++) {
                dp[i][j] = -1;
            }
        }

        int startPos = -1;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == k) {
                startPos = i;
                break;
            }
        }

        if (startPos == -1) return 0;

        dp[startPos][startPos] = 0;
        for (int i = startPos; i >= 0; i--) {
            for (int j = startPos; j <= nums.length - 1; j++) {
                if (i - 1 >= 0) {
                    dp[i - 1][j] = dp[i][j] + (nums[i - 1] > k ? 1 : -1);

                }
                if (j + 1 <= nums.length - 1) {
                    dp[i][j + 1] = dp[i][j] + (nums[j + 1] > k ? 1 : -1);
                }
            }
        }


        int res = 0;
        for (int i = startPos; i >= 0; i--) {
            for (int j = startPos; j <= nums.length - 1; j++) {
                if (dp[i][j] == 0 || dp[i][j] == 1) {
                    res++;
                }
            }
        }

        return res;
    }

    public int countSubarrays2(int[] nums, int k) {
        int pos = 0, n = nums.length;
        while (nums[pos] != k) ++pos;

        var cnt = new HashMap<Integer, Integer>();
        // i=pos 的时候 x 是 0，直接记到 cnt 中，这样下面不是大于 k 就是小于 k
        cnt.put(0, 1);
        for (int i = pos - 1, x = 0; i >= 0; --i) { // 从 pos-1 开始累加 x
            x += nums[i] < k ? 1 : -1;
            cnt.merge(x, 1, Integer::sum);
        }

        // i=pos 的时候 x 是 0，直接加到答案中，这样下面不是大于 k 就是小于 k
        int ans = cnt.get(0) + cnt.getOrDefault(-1, 0);
        for (int i = pos + 1, x = 0; i < n; ++i) { // 从 pos+1 开始累加 x
            x += nums[i] > k ? 1 : -1;
            ans += cnt.getOrDefault(x, 0) + cnt.getOrDefault(x - 1, 0);
        }
        return ans;
    }
}
